''' 
本程序为参照《地球物理中的有限单元法》3.3.1章节练习内容而编写，适用于
第一类边界条件、三角元剖分、线性插值的位场延拓
author:taqikema
date:20250803
history:
    20250803,初次编写
'''

import numpy as np

def init_raw_data():
    """返回 nd, node_no, node_xy, node1_no_value
    """
    nd = 15     # 节点总数
    node_no = [
    [1, 4, 3],
    [5, 2, 4],
    [5, 3, 2],
    [5, 6, 3],
    [5, 4, 7],
    [5, 7, 8],
    [5, 8, 6],
    [8, 9, 6],
    [11, 7, 10],
    [8, 7, 11],
    [8, 11, 12],
    [8, 12, 9],
    [10, 13, 14],
    [11, 10, 14],
    [11, 14, 15],
    [11, 15, 12]
    ]
    node_no = np.array(node_no)     # 节点编号
    
    node_xy = [
    [0, 0],
    [0, 1],
    [0, 2],
    [1, 0.5],
    [1, 1.2],
    [1, 2],
    [2, 1],
    [2, 1.5],
    [2, 2],
    [3, 0.7],
    [3, 1.3],
    [3, 2],
    [4, 0.5],
    [4, 1.2],
    [4, 2]
    ]
    node_xy = np.array(node_xy)     # 节点坐标
    
    node1_no_value = [
        [1, 1],
        [2, 0.5],
        [3, 0.33],
        [4, 0.46],
        [6, 0.3],
        [7, 0.25],
        [9, 0.23],
        [10, 0.14],
        [12, 0.17],
        [13, 0.08],
        [14, 0.11],
        [15, 0.12],
    ]
    node1_no_value = np.array(node1_no_value)   # 第一类边界的节点号和场值
    
    return nd, node_no, node_xy, node1_no_value

def calc_half_band_width(node_no):
    """
    计算半带宽值，先计算每个单元中三个节点间插值的最大值,再选出所有单元中的最大差值
    """
    # 每行最大元素减最小元素即为该行最大绝对差值
    row_ranges = np.ptp(node_no, axis=1)  # ptp = peak to peak (max - min)
    return np.max(row_ranges)

def assemble_unit_coefficient_matrix(unit_node_xy):
    """
    计算单个单元的系数矩阵
    参数:
        unit_node_xy : numpy.ndarray, 形状为 (3, 2)
    
    返回:
        unit_coefficient_matrix : numpy.ndarray, 形状为 (3, 3)
    """
    A = np.array([unit_node_xy[1,1] - unit_node_xy[2,1], unit_node_xy[2,1] - unit_node_xy[0,1], unit_node_xy[0,1] - unit_node_xy[1,1]])
    B = np.array([unit_node_xy[2,0] - unit_node_xy[1,0], unit_node_xy[0,0] - unit_node_xy[2,0], unit_node_xy[1,0] - unit_node_xy[0,0]])
    S = 2.0 * (A[0]*B[1] - A[1]*B[0])
    return (A.T * A + B.T * B) / S
    

def assemble_coefficient_matrix(nd, node_no, node_xy, half_band_width):
    """
    组装半带宽总体系数矩阵
    """
    K = np.zeros((nd, nd))
    # 先计算完整的 K矩阵
    for i in range(np.size(node_no, 1)):
        node_index = node_no[i,:]
        unit_coefficient_matrix = assemble_unit_coefficient_matrix(node_xy[node_index, :])
        K[node_index,node_index] += unit_coefficient_matrix
        
    # 再转换成 SK矩阵
    SK = np.zeros((nd, half_band_width))
    for j in range(nd):
        for k in range(j):
            SK[j,half_band_width-j+k] = K[j,k]
    
    return SK


def use_bc_Dirichlet(node1_no_value, SK):
    """
    在定带宽总体系数矩阵和右侧向量上加上第一类边界条件的节点值
    """
    node1_value = np.zeros((1, np.size(node1_no_value, 1)))
    SK[node1_no_value[:,0], -1] *= 1.0E10
    node1_value = SK[node1_no_value[:,0], -1] * node1_no_value[:,1]
    
    return SK, node1_value

def solve_LES(SK, node1_value):
    """
    在
    """
    return x

def potential_field_continuation():
    """主函数, 位场延拓, 第一类边界条件、三角元剖分、线性插值"""
    print("=== 位场延拓, 第一类边界条件、三角元剖分、线性插值 ===")
    
    # 输入原始数据
    nd, node_no, node_xy, node1_no_value = init_raw_data()
    
    # 计算半带宽
    half_band_width = calc_half_band_width(node_no)
    
    # 组装半带宽总体系数矩阵
    SK  = assemble_coefficient_matrix(nd, node_no, node_xy, half_band_width)
    
    # 代入第一类边界条件
    SK, node1_value = use_bc_Dirichlet(node1_no_value, SK)
    
    # 求解定带宽储存的对称带型线性方程组
    node_value = solve_LES(SK, node1_value)
    
    # 测试编号 5/8/11的场值计算是否正确
    assert(node_value[[4,7,10]], [0.376, 0.244, 0.162])

if __name__ == "__main__":
    potential_field_continuation()
